import numpy as np
import pandas as pd
from efficient_apriori import apriori
from mlxtend.frequent_patterns import apriori
from mlxtend.preprocessing import TransactionEncoder
from mlxtend.frequent_patterns import association_rules


def apply_price(x):
    if x < 20:
        return 'Very Low'
    elif x < 35:
        return 'Low'
    elif x < 50:
        return 'General'
    elif x < 65:
        return 'High'
    else:
        return 'Very High'


def apply_points(x):
    if x <= 85:
        return 'Level 1'
    elif x <= 90:
        return 'Level 2'
    elif x <= 95:
        return 'Level 3'
    else:
        return 'Level 4'


if __name__ == '__main__':
    pd.set_option('display.expand_frame_repr', False)
    pd.set_option('display.max_columns', None)
    pd.set_option('display.max_rows', None)

    df = pd.DataFrame(pd.read_csv('winemag-data_first150k.csv'))
    pp_df = df[['price', 'points']]
    pp_df['price'] = pp_df['price'].apply(apply_price)
    pp_df['points'] = pp_df['points'].apply(apply_points)
    pp_array = np.array(pp_df.head(10000))
    pp_list = pp_array.tolist()

    te = TransactionEncoder()  # 定义模型
    df_tf = te.fit_transform(pp_list)
    df = pd.DataFrame(df_tf, columns=te.columns_)
    frequent_itemsets = apriori(df, min_support=0.003, use_colnames=True)  # use_colnames=True表示使用元素名字，默认的False使用列名代表元素
    # frequent_itemsets = apriori(df,min_support=0.05)
    frequent_itemsets.sort_values(by='support', ascending=False, inplace=True)  # 频繁项集可以按支持度排序
    print(frequent_itemsets[frequent_itemsets.itemsets.apply(lambda x: len(x)) >= 2])  # 选择长度 >=2 的频繁项集
    association_rule = association_rules(frequent_itemsets, metric='confidence',
                                         min_threshold=0.5)  # metric可以有很多的度量选项，返回的表列名都可以作为参数
    association_rule.sort_values(by='leverage', ascending=False, inplace=True)  # 关联规则可以按leverage排序
    print(association_rule)